/*
 * Copyright(C) 2021. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "Yolov5PostProcess.h"
#include "MxBase/Log/Log.h"
#include "MxBase/Maths/FastMath.h"
#include "MxBase/CV/ObjectDetection/Nms/Nms.h"



namespace {
    const int MODEL_INPUT_SIZE = 640;
    const int ALIGN_LEFT = 16;
    const int CONFIDENCE_IDX = 4;
    const int LABEL_START_OFFSET = 5;
    const int OUTPUT_DIMS = 3;
    const int XOFFSET = 2;
    const int YOFFSET = 3;
    const int AVG_PARAM = 2;
    const float EPSILON = 1e-6;
}


Yolov5PostProcess &Yolov5PostProcess::operator=(const Yolov5PostProcess &other) {
    if (this == &other) {
        return *this;
    }
    ObjectPostProcessBase::operator=(other);
    objectnessThresh_ = other.objectnessThresh_; // Threshold of objectness value
    iouThresh_ = other.iouThresh_;
    paddingType_ = other.paddingType_;

    return *this;
}

APP_ERROR Yolov5PostProcess::Init(const std::map<std::string, std::shared_ptr<void>> &postConfig) 
{
    LogDebug << "Start to Init Yolov5PostProcess.";
    APP_ERROR ret = ObjectPostProcessBase::Init(postConfig);
    if (ret != APP_ERR_OK) {
        LogError << GetError(ret) << "Fail to superInit in ObjectPostProcessBase.";
        return ret;
    }

    LogDebug << "End to Init Yolov5PostProcess.";
    return APP_ERR_OK;
}

APP_ERROR Yolov5PostProcess::DeInit() {
    return APP_ERR_OK;
}

void Yolov5PostProcess::ConstructBoxFromOutput(float *output, size_t offset, std::vector<MxBase::ObjectInfo> &objectInfo,
    const MxBase::ResizedImageInfo &resizedImageInfo)
{
    size_t index = offset * (classNum_ + LABEL_START_OFFSET);
    if (output[index + CONFIDENCE_IDX] <= objectnessThresh_) {
        return;
    }
    int classId = -1;
    float maxProb = scoreThresh_;
    for (uint32_t j = LABEL_START_OFFSET; j < classNum_ + LABEL_START_OFFSET; j++) {
        if (output[index + j] * output[index + CONFIDENCE_IDX] > maxProb) {
            maxProb = output[index + j] * output[index + CONFIDENCE_IDX];
            classId = j - LABEL_START_OFFSET;
        }
    }
    if (classId < 0) {
        return;
    }
    int tmpResizedWidth = resizedImageInfo.widthResize;
    int tmpResizedHeight = resizedImageInfo.heightResize;
    double division = 1;
    // if (std::fabs(resizedImageInfo.keepAspectRatioScaling) > EPSILON) {
    //     division = resizedImageInfo.keepAspectRatioScaling;
    // }
    if (tmpResizedWidth == tmpResizedHeight && tmpResizedHeight == MODEL_INPUT_SIZE) {
        tmpResizedWidth = std::round(resizedImageInfo.widthOriginal * division);
        tmpResizedHeight = std::round(resizedImageInfo.heightOriginal * division);
    }
    // int offsetLeft = (MODEL_INPUT_SIZE - tmpResizedWidth) / AVG_PARAM;
    // int offsetTop = (MODEL_INPUT_SIZE - tmpResizedHeight) / AVG_PARAM;

    // if (paddingType_ == 0) {
    //     offsetTop = offsetTop % AVG_PARAM == 0 ? offsetTop : offsetTop - 1;
    //     offsetLeft = offsetLeft < ALIGN_LEFT ? 0 : offsetLeft / ALIGN_LEFT * ALIGN_LEFT;
    // }
    // auto leftX = (output[index] - output[index + XOFFSET] / AVG_PARAM - offsetLeft) / division;
    // auto leftY = (output[index + 1] - output[index + YOFFSET] / AVG_PARAM - offsetTop) / division;
    // auto rightX = (output[index] + output[index + XOFFSET] / AVG_PARAM - offsetLeft) / division;
    // auto rightY = (output[index + 1] + output[index + YOFFSET] / AVG_PARAM - offsetTop) / division;

    double scale_ratio_w = (double)tmpResizedWidth / MODEL_INPUT_SIZE;
    double scale_ratio_h = (double)tmpResizedHeight / MODEL_INPUT_SIZE;

    auto leftX = (output[index] - output[index + XOFFSET] / AVG_PARAM) * scale_ratio_w;
    auto leftY = (output[index + 1] - output[index + YOFFSET] / AVG_PARAM) * scale_ratio_h;
    auto rightX = (output[index] + output[index + XOFFSET] / AVG_PARAM) * scale_ratio_w;
    auto rightY = (output[index + 1] + output[index + YOFFSET] / AVG_PARAM) * scale_ratio_h;

    MxBase::ObjectInfo obj;
    obj.x0 = leftX < 0.0 ? 0.0 : leftX;
    obj.y0 = leftY < 0.0 ? 0.0 : leftY;
    obj.x1 = rightX > resizedImageInfo.widthOriginal ? resizedImageInfo.widthOriginal : rightX;
    obj.y1 = rightY > resizedImageInfo.heightOriginal ? resizedImageInfo.heightOriginal : rightY;
    obj.confidence = maxProb;
    obj.classId = classId;
    obj.className = configData_.GetClassName(obj.classId);
    if (maxProb < separateScoreThresh_[obj.classId])
        return;
    objectInfo.push_back(obj);
}

// 将处理好的推理结果放入ObjectInfo
void Yolov5PostProcess::ObjectDetectionOutput(const std::vector<MxBase::TensorBase> &tensors,
                                              std::vector<std::vector<MxBase::ObjectInfo>> &objectInfos,
                                              const std::vector<MxBase::ResizedImageInfo> &resizedImageInfos) {
    LogDebug << "Yolov5PostProcess start to write results.";
    if (tensors.size() == 0) {
        return;
    }

    uint32_t batchSize = tensors[0].GetShape()[0];
    size_t rows = tensors[0].GetSize() / ((classNum_ + LABEL_START_OFFSET) * batchSize);
    for (size_t k = 0; k < batchSize; k++) {
        auto output = (float*)GetBuffer(tensors[0], k);
        std::vector<MxBase::ObjectInfo> objectInfo;
        for (size_t i = 0; i < rows; ++i) {
            ConstructBoxFromOutput(output, i, objectInfo, resizedImageInfos[k]);
        }
        MxBase::NmsSort(objectInfo, iouThresh_);
        objectInfos.push_back(objectInfo);
    }
    LogDebug << "Yolov5PostProcess write results successed.";
}

APP_ERROR Yolov5PostProcess::Process(const std::vector<MxBase::TensorBase> &tensors,
                                     std::vector<std::vector<MxBase::ObjectInfo>> &objectInfos,
                                     const std::vector<MxBase::ResizedImageInfo> &resizedImageInfos,
                                     const std::map<std::string, std::shared_ptr<void>> &paramMap) {
    LogDebug << "Start to Process Yolov5PostProcess.";
    APP_ERROR ret = APP_ERR_OK;
    if (resizedImageInfos.size() == 0) {
        ret = APP_ERR_INPUT_NOT_MATCH;
        LogError << GetError(ret) << "resizedImageInfos is not provided which is necessary for Yolov5PostProcess.";
        return ret;
    }
    auto inputs = tensors;
    ret = CheckAndMoveTensors(inputs);
    if (ret != APP_ERR_OK) {
        LogError << GetError(ret) << "CheckAndMoveTensors failed.";
        return ret;
    }
    
    ObjectDetectionOutput(inputs, objectInfos, resizedImageInfos);

    // LogObjectInfos(objectInfos);
    LogDebug << "End to Process Yolov5PostProcess.";
    return APP_ERR_OK;
}
